from template import *
from collect import collect

import sys

itime = -1
if len(sys.argv) == 2:
	t = int(sys.argv[1])
elif len(sys.argv) > 2:
	print("error: should have one argument, as time step")
else:
	t = itime



x_step = 20

amplification_factor = 80.0


data_temp = collect("/Diagnostic/", "radiative_energy_collision")
nx = data_temp.shape[3]


dx = 0.5e-5  # unit (m)
x = np.linspace(0, nx * dx, nx)
x = x * amplification_factor
x_less = x[::x_step]


xmin = x.min()
xmax = x.max()

x0 = int(nx / 2)
x1 = 100


list_label = [r"$\mathrm{Colomb}$", r"$\mathrm{Ion\ D}$", r"$\mathrm{Exc\ D}$", r"$\mathrm{Ion\ C}$", r"$\mathrm{Ion\ C^{1+}}$", r"$\mathrm{Ion\ C^{2+}}$", r"$\mathrm{Exc\ C}$", r"$\mathrm{Exc\ C^{1+}}$", r"$\mathrm{Exc\ C^{2+}}$", r"$\mathrm{Exc\ C^{3+}}$"]


##inite the fig of matplotlib
fig=plt.figure(figsize=(10,8))
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

##==================================================================
ax0=fig.add_subplot(1,1,1)

val = collect("/Diagnostic/", "radiative_energy_collision")
val = val[:]
val_2d = val[t,0,:,:]
n_collide = val_2d.shape[0]

radiative_energy_collision_sum = np.zeros(val_2d.shape[1])

for i_collide in np.arange(1, n_collide):
	#val_1d = np.transpose(val[t, 0, 0, :])
	#print( "Electron density: ",val_1d[x0], val_1d[x1] )
	val_1d = val_2d[i_collide,:]
	line0=ax0.plot(x, val_1d, label = list_label[i_collide])

	radiative_energy_collision_sum = radiative_energy_collision_sum + val_1d

line0=ax0.plot(x, radiative_energy_collision_sum, label = "Total")

radiative_energy_flux = 0.0
for i in range(radiative_energy_collision_sum.shape[0]):
	radiative_energy_flux = radiative_energy_flux + radiative_energy_collision_sum[i] * dx
print("radiative_energy_flux: ", radiative_energy_flux)


#ax0.grid(True)
ax0.legend(framealpha = 0.2)
ax0.set_xlim((xmin, xmax))
#ax0.set_ylim((0.0, 5.1))
ax0.set_xlabel(r"$x \mathrm{(m)}$", fontsize = label_fontsize)
ax0.set_ylabel(r"$P \mathrm{_{loss}(Wm^{-3})}$", fontsize = label_fontsize)
ax0.set_title(r"$n \mathrm{=1.0\times 10^{18} m^{-3}}$")

#major_ticks = np.arange(0, 5.1, 1.5)
#ax0.set_yticks(major_ticks)
#ax0.set_ylabel(r"$n\ \mathrm{(10^{19}m^{-3})}$", fontsize = label_fontsize)

ax0.annotate(r"$\mathbf{(a)}$", xy=get_axis_limits(ax0), annotation_clip=False)



pdf_file_name = "radiative_energy_collision" + str(t) + ".svg"
fig.savefig(pdf_file_name, dpi = 300)
##fig.show()       #when the program finishes,the figure disappears
#plt.axis('equal')
#plt.show()         #The command is OK



#heat flux
val = collect("/Diagnostic/", "heat_flux_left")
val = val[:]
print("heat_flux_left: ", val[t,0,0,0] + val[t,0,0,1], val[t,0,0,:])

val = collect("/Diagnostic/", "heat_flux_right")
val = val[:]
print("heat_flux_right: ", val[t,0,0,0] + val[t,0,0,1], val[t,0,0,:])